from collections import OrderedDict

import torch
import torch.nn as nn
from detection.layers import FrozenBatchNorm2d
from torchvision import models


class VGG(nn.Module):
    def __init__(self, vgg, cfg):
        super().__init__()
        self.backbone = vgg.features[:-1]
        self.ada_layers = cfg.ADV.LAYERS

    def forward(self, x):
        adaptation_feats = []
        idx = 0
        for i in range(14):
            x = self.backbone[i](x)
        if self.ada_layers[idx]:
            adaptation_feats.append(x)

        idx += 1
        for i in range(14, 21):
            x = self.backbone[i](x)
        if self.ada_layers[idx]:
            adaptation_feats.append(x)

        idx += 1
        for i in range(21, len(self.backbone)):
            x = self.backbone[i](x)
        if self.ada_layers[idx]:
            adaptation_feats.append(x)

        return [x], adaptation_feats




def vgg(cfg, pretrained=True):
    backbone_name = cfg.MODEL.BACKBONE.NAME
    vgg = models.vgg.__dict__[backbone_name](pretrained=pretrained)
    vgg = VGG(vgg, cfg)
    vgg.out_channels = 512

    for layer in range(10):
        for param in vgg.backbone[layer].parameters():
            param.requires_grad = False

    return vgg


